from multiprocessing import (Pool, Queue, Process, Value, Array, Lock, cpu_count, Event)
from sage.combinat.backtrack import SearchForest
from sage.rings.integer import Integer
from sage.misc.abstract_method import abstract_method
from sage.categories.finite_enumerated_sets import FiniteEnumeratedSets
from sage.structure.unique_representation import UniqueRepresentation
import threading
import time

class threadPerso(threading.Thread):
	def __init__(self, ob, fin, nom = ''):
		threading.Thread.__init__(self)
		self._ob = ob
		self._nom = nom
		self._fin = fin
		self._stopevent = threading.Event()
	
	def run(self):
		r"""Commentaire """
		while (not self._stopevent.isSet() and self._ob._qres.qsize() > 0) or self._fin.value == False:
			self._ob._res = self._ob.reduce_function(self._ob._res, self._ob._qres.get())

	def stop(self):
		self._stopevent.set( )


class SearchForest2(SearchForest):
	r""" Commentaire """

	def __init__(self, roots = None, children = None, post_process = None,
		    statistic = None, reduce_function = None, reduce_init = None, algorithm = 'depth', category = None ):
	
		SearchForest.__init__(self,roots,children,post_process, algorithm, category)
		""" Initialisation du tunel utilise pour l'envoie des resultats
		self._parent_conn, self._child_conn = Pipe()
		"""
		
		""" Initialisation du post_process """
		if post_process is None:
			self.post_process = lambda x: x
		if statistic is not None :
			self.statistic = statistic
		if reduce_function is not None :
			self.reduce_function = reduce_function
		self._reduce_init = reduce_init
		self._res = None

	@abstract_method
	def statistic(self, x):
		r"""Commentaire"""
		

	@abstract_method
	def reduce_function(self, res):
		r"""Commentaire """

	def traitement(self, i, l,a, fin):
		if type(self._reduce_init) == list:
			res = list(self._reduce_init)
		else:
			res = self._reduce_init
	
		while fin.value == False:			#Boucle tant que tous les processes n'ont pas finie de travaille.
			while self._q.qsize() > 0:		#Boucle tant qu'il y a des elements dans la queue.
				a[i] = 0			#Tableau qui pour l'indice du processe indique qu'il travaille 0.
				try:
					o = self._q.get(False)	# On lit dans le queue mais pas en bloquant d'ou le False. Si rien a lire exception.
					for ch in self.children(o): # Ajout des fils dans la queue
						self._q.put(ch)
				
					if self.post_process(o) is not None :   # Si post_process return un objet different de none alors
						res = self.reduce_function(res, self.statistic(o)) # Il enregistre l'objet
				except:
					continue 				# Exception leve lorsqu'on lit dans une queue vide
			a[i] = 1						# Indique que le processes n'a pas trouve d'element dans la queue.
		#l.acquire()
		self._qres.put(res)
		#l.release()
		
	def run(self):
		r"""Commentaire """
		
		if self._res is None:
			self._q = Queue()			# Queue pour la traverse de l'arbre.
			self._qres = Queue()			# Queue pour recuperer le resultat de chaque processe
			l = Lock()				# Verrou non utilise pour le moment
			fin = Value('b',False)			# Variable qui permet de dire a tout les processes que le travaille est finie
			
			arr = Array('i',range(cpu_count()))	# Tableau pour connaitre l'etat de chaque processe
			for i in range(cpu_count()):
				arr[i] = 0			# Initialisation du tableau
			affiche = threadPerso(self, fin, 'Affiche') # Thread de recuperation des resultats a la fin de chaque processe
		
			if type(self._reduce_init) == list:	# Il faut traite le cas des listes.
				self._res = list(self._reduce_init)
			else:
				self._res = self._reduce_init

			for i in self.roots():			# Initialisation de la queue
				self._q.put(i)
			
			processes = [Process(target = self.traitement, args = (i,l,arr,fin)) for i in range(cpu_count())] # Creation des processe

			affiche.start()				# Demarage du Thread de recuperation
			for p in processes:
				p.start()			# Demarage des processes
				
			while fin.value == False :		# Attend que tout les processes est finie les travailles pour les arretes
				cp = 0
				for i in range(cpu_count()):
					if arr[i] == 1:
						cp = cp + 1
				if cp == cpu_count():
					fin.value = True

			for p in processes:
				p.join()
			
			affiche.join()
			#self._qres.close()
			#self._q.close()

		return self._res
	
	def initRes(self):					# Reinitialisation de res. Afin de refaire le calcul
		if self._res != None :
			self._res = None
			
S = SearchForest2( [[]],                            # root
    lambda l: [l+[0], l+[1]] if len(l) < 12 else [], # children
    lambda x: x,                                    # post-process
    lambda x: var('x')**len(x),                     # statistic
    lambda x,y: x+y,                                # reduce_function
    0)                                              # reduce_init
"""
S.run()
"""